import socket
from random import choice
from time import sleep
from copy import deepcopy

import sys
import math
import random
import numpy as np

from alphazero.MCTS import MCTS
from alphazero.hex.HexGame import HexGame
from alphazero.hex.HexPlayers import *
from alphazero.hex.pytorch.NNet import NNetWrapper as NNet

from alphazero.utils import *

class MctsAgent():
    """This class describes the default Hex agent. It will randomly send a
    valid move at each turn, and it will choose to swap with a 50% chance.

    agent only maintain the state of board, the decision will made by mcts
    """

    HOST = "127.0.0.1"
    PORT = 1234

    def run(self):
        """A finite-state machine that cycles through waiting for input
        and sending moves.
        """

        self._board_size = 0
        self._board = []  # record current state of board
        self._colour = ""
        self._turn_count = 1

        states = {
            1: MctsAgent._connect,
            2: MctsAgent._wait_start,
            3: MctsAgent._make_move,
            4: MctsAgent._wait_message,
            5: MctsAgent._close
        }

        res = states[1](self)
        while (res != 0):
            res = states[res](self)

    def _connect(self):
        """Connects to the socket and jumps to waiting for the start
        message.
        """

        self._s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._s.connect((MctsAgent.HOST, MctsAgent.PORT))

        return 2

    def _wait_start(self):
        """Initialises itself when receiving the start message, then
        answers if it is Red or waits if it is Blue.
        """

        data = self._s.recv(1024).decode("utf-8").strip().split(";")
        if (data[0] == "START"):
            self._board_size = int(data[1])
            self._colour = data[2]
            self._board_str = ",".join(["0" * self._board_size] *
                                       self._board_size)
            
            # initialize nnet player
            self._game = HexGame(self._board_size)
            nnet = NNet(self._game)
            nnet.load_checkpoint('agents/Group888/alphazero/pretrained_models', 'best.pth.tar')
            args = dotdict({'numMCTSSims': 400, 'cpuct': 1.0})
            self._mcts = MCTS(self._game, nnet, args)
            self._player = lambda x: np.argmax(self._mcts.getActionProb(x, temp=0))

            if (self._colour == "R"):
                return 3
            else:
                return 4

        else:
            print("ERROR: No START message received.")
            return 0

    def _make_move(self):
        """Makes a random valid move. It will choose to swap with
        a coinflip.
        """
        board = self._get_board_from_str()
        side = 1 if self._colour == "R" else -1
        canonicalBoard = self._game.getCanonicalForm(board, side)

        if self._turn_count == 2 and self.isSwap():
            msg = "SWAP\n"
        else:
            action = self._player(canonicalBoard)
            if side == -1:
                action = self._game.remapAction(action)
            move = (action // self._game.n, action % self._game.n)
            msg = f"{move[0]},{move[1]}\n"

        self._s.sendall(bytes(msg, "utf-8"))

        return 4

    def isSwap(self):
        low_bound = int(self._board_size * 0.25)
        up_bound = int(self._board_size * 0.75)
        if self._last_move[0] < low_bound or self._last_move[0] > up_bound:
            return False
        if self._last_move[1] < low_bound or self._last_move[1] > up_bound:
            return False
        return True

    def _wait_message(self):
        """Waits for a new change message when it is not its turn."""

        self._turn_count += 1

        data = self._s.recv(1024).decode("utf-8").strip().split(";")
        if (data[0] == "END" or data[-1] == "END"):
            return 5
        else:
            # CHANGE;MOVE;BOARD;COLOUR

            if (data[1] == "SWAP"):
                self._colour = self.opp_colour()
            else:
                x, y = data[1].split(",")
                self._last_move = (int(x), int(y))
                self._board_str = data[2]

            if (data[-1] == self._colour):
                return 3

        return 4

    def _close(self):
        """Closes the socket."""

        self._s.close()
        return 0

    def _get_board_from_str(self):
        board = np.ndarray((self._board_size, self._board_size), dtype=int)
        ss = self._board_str.split(',')
        for i, s in enumerate(ss):
            for j, c in enumerate(s):
                if c == "R":
                    board[i][j] = 1
                elif c == "B":
                    board[i][j] = -1
                else:
                    board[i][j] = 0
        return board

    def opp_colour(self):
        """Returns the char representation of the colour opposite to the
        current one.
        """

        if self._colour == "R":
            return "B"
        elif self._colour == "B":
            return "R"
        else:
            return "None"


if (__name__ == "__main__"):
    agent = MctsAgent()
    agent.run()
